import os
import torch
import torch.nn as nn
import argparse
from torch.backends import cudnn
from PIL import Image
import random, string
from torchvision import transforms
import numpy as np
from torchvision.utils import save_image
from network.noise_layers.stargan.model import Generator


class StarGAN(nn.Module):

    def __init__(self, c_trg=3, image_size=128, target="/home/fang_guotong/datasets/watermark_dataset/CelebA-HQ128/test_128/", temp="temp/"):
        super(StarGAN, self).__init__()
        self.c_trg = c_trg
        self.image_size = image_size
        self.temp = temp
        self.target = target
        self.attrs = ['Black_Hair', 'Blond_Hair', 'Brown_Hair', 'Male', 'Young']
        self.G = Generator(64, 5, 6)
        G_path = os.path.join('network/noise_layers/stargan', str(self.image_size),'200000-G.ckpt')
        self.G.load_state_dict(torch.load(G_path, map_location=f'cpu'))
        # print(self.G.state_dict().keys())
        if not os.path.exists(temp): os.mkdir(temp)
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
        ])

    def get_temp(self):
        return self.temp + ''.join(random.sample(string.ascii_letters + string.digits, 16)) + ".png"  # png or jpg?

    def get_target(self):
        idx = random.randint(162771, 182637)
        return self.target + str(idx).zfill(6) + ".png"

    def denorm(self, x):
        """Convert the range from [-1, 1] to [0, 1]."""
        out = (x + 1) / 2
        return out.clamp_(0, 1)

    def forward(self, image_cover_mask):
        image, cover_image, mask = image_cover_mask
        device = image.device
        self.G = self.G.to(device)
        self.G.eval()

        noised_image = torch.zeros_like(image)

        for i in range(image.shape[0]):
            single_image = ((image[i].clamp(-1, 1).permute(1, 2, 0) + 1) / 2 * 255).add(0.5).clamp(0, 255).to('cpu', torch.uint8).numpy()
            im = Image.fromarray(single_image)
            file = self.get_temp()
            while os.path.exists(file):
                file = self.get_temp()
            try:
                self.test(self.get_target(), im, mask[i], file, device)
                fake = np.array(Image.open(file), dtype=np.uint8)
                os.remove(file)
                noised_image[i] = self.transform(fake).to(device)
            except Exception as e:
                print(f"[ERROR] StarGAN error on image {i}: {e}")
                noised_image[i] = image[i]
                raise e
        return noised_image


    def test(self, face_image_path, body_image, label, output_path, device):
        regular_image_transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
        ])

        with torch.no_grad():
            x_real = regular_image_transform(body_image).unsqueeze(0).to(device)
            c_org = label.to(device)
            c_trg_list = self.create_labels(c_org, 5, selected_attrs=self.attrs, device=device)
            if self.c_trg is None:
                self.c_trg = random.randint(0, 4)
            x_fake = self.G(x_real, c_trg_list[self.c_trg].unsqueeze(0))
            save_image(self.denorm(x_fake), output_path)


    def create_labels(self, c_org, c_dim=5, dataset='CelebA', selected_attrs=None, device=None):
        if device is None:
            device = c_org.device

        hair_color_indices = []
        if dataset == 'CelebA':
            for i, attr_name in enumerate(selected_attrs):
                if attr_name in ['Black_Hair', 'Blond_Hair', 'Brown_Hair', 'Gray_Hair']:
                    hair_color_indices.append(i)

        c_trg_list = []
        for i in range(c_dim):
            if dataset == 'CelebA':
                c_trg = c_org.clone()
                if i in hair_color_indices:
                    c_trg[i] = 1
                    for j in hair_color_indices:
                        if j != i:
                            c_trg[j] = 0
                else:
                    c_trg[i] = (c_trg[i] == 0)
            elif dataset == 'RaFD':
                c_trg = self.label2onehot(torch.ones(c_org.size(0)) * i, c_dim)
            c_trg_list.append(c_trg.to(device))
        return c_trg_list

    
    def print_weights_structure(self, weight_path):
        """
        加载权重文件并打印其包含的所有键名，展示权重的模型结构
        """
        try:
            # 仅加载权重文件，不加载到模型
            weights = torch.load(weight_path, map_location='cpu')  # 用CPU加载避免CUDA问题
            # print(f"\n===== 权重文件 ({os.path.basename(weight_path)}) 结构 =====")
            # print("权重包含的键名列表：")
            # for idx, key in enumerate(weights.keys()):
            #     print(f"[{idx}] {key}")  # 打印每个键名（即权重对应的层名）
            # print(f"=========================================\n")
            return weights
        except Exception as e:
            print(f"加载权重文件失败：{e}")
            return None
        
    def compare_model_and_weights(self, weight_path):
        """
        对比当前模型的层名称与权重文件的键名，找出不匹配的部分
        """
        weights = self.print_weights_structure(weight_path)
        if weights is None:
            return
        
        # 获取模型的所有参数名称
        model_param_names = [name for name, _ in self.G.named_parameters()]
        # 获取权重的所有键名
        weight_keys = list(weights.keys())
        
        # 找出权重中有但模型中没有的键名（对应你遇到的Unexpected key错误）
        unexpected_keys = [key for key in weight_keys if key not in model_param_names]
        # 找出模型中有但权重中没有的键名（可能导致参数未初始化）
        missing_keys = [name for name in model_param_names if name not in weight_keys]
        
        if unexpected_keys:
            print("===== 权重中存在但模型中没有的键名（Unexpected keys） =====")
            for key in unexpected_keys:
                print(key)
        
        if missing_keys:
            print("\n===== 模型中存在但权重中没有的键名（Missing keys） =====")
            for name in missing_keys:
                print(name)
        
        if not unexpected_keys and not missing_keys:
            print("\n模型结构与权重结构完全匹配！")